#! /usr/bin/env python

from worker import Worker, WorkerError
import select
import os
from time import sleep
import signal

class ContractorError(Exception):
    """Base class for all exceptions raised by this module.
    """
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return str(self.value)

class Contractor:
    def __init__(self, number, program, stderr = '/dev/null'):
        self._check_mode(program, stderr)

        self.program    = program
        self.stderr     = stderr
        
        self.number     = number
        self.num_child  = number
        # pid: worker
        self.employee   = {}
        self.fd2pid     = {}
        self.idle       = []
        self.busy       = []

        signal.signal(signal.SIGCHLD, self.sigchld_handler())

    def _check_mode(self, program, stderr):
        if os.path.abspath(program) != program:
            raise ContractorError('%s: It is not absolute path' % program)
        if not os.path.exists(program):
            raise ContractorError('%s: No such file' % program)
        if not os.access(program, os.R_OK|os.X_OK):
            raise ContractorError('%s: Permission denied' % program)
        if not os.access(stderr, os.W_OK):
            raise ContractorError('%s: Permission denied' % stderr)

    def recruit(self):
        for i in xrange(self.num_child):
            worker = Worker(self.program, self.stderr)
            worker.run()

            pid = worker.pid
            fd  = worker.outfd

            self.employee[pid]  = worker

            self.fd2pid[fd]     = pid
            self.idle.append(pid)

            self.num_child -= 1

    def sigchld_handler(self):
        def _unname_(signum, frame):
            self.num_child += 1
        return _unname_

    def trial(self, time = 1):
        """Wait a few secs, during this moment, 
        any child abort means worker has a fatal fault.
        """
        sleep(time)
        if self.num_child != 0:
            raise ContractorError('%s: Some workers on strike' % self.program)

    def handle_resign(self):
        for i in xrange(self.num_child):
            pid, ecode = os.wait()
            fd = self.employee[pid].outfd

            try:
                self.idle.remove(pid)
            except ValueError:
                pass
            try:
                self.busy.remove(pid)
            except ValueError:
                pass
            self.fd2pid.pop(fd)
            self.employee.pop(pid)

    def idle_count(self):
        return len(self.idle)

    def assign_task(self, id, task):
        try:
            pid = self.idle.pop()
            worker = self.employee[pid]
            worker.request(id, task)
            self.busy.append(pid)
        except IndexError:
            raise ContractorError("all workers are busy now")

    def wait_response(self, timeout = -1):
        epoll = select.epoll()
        for pid in self.busy:
            fd = self.employee[pid].outfd
            epoll.register(fd, select.EPOLLIN)

        results = []
        try:
            readable = epoll.poll(timeout)
            complete = [self.fd2pid[fd] for fd, event in readable]
            for pid in complete:
                worker = self.employee[pid]
                id, result = worker.response()
                self.busy.remove(pid)
                self.idle.append(pid)
                results.append((id, result))
        except IOError:
            pass

        return results

    def quit_graceful(self):
        ret = []
        while len(self.idle) != self.number:
            ret.extend(self.wait_response())
        self.quit()
        return ret

    def quit(self):
        for pid, worker in self.employee.items():
            fd = worker.outfd
            try:
                self.idle.remove(pid)
            except ValueError:
                pass
            try:
                self.busy.remove(pid)
            except ValueError:
                pass
            self.fd2pid.pop(fd)
        self.employee = {}
        signal.signal(signal.SIGCHLD, signal.SIG_DFL)

if __name__ == '__main__':
    def generator():
        id = 1
        while True:
            yield id
            id += 1

    IDGenerator = generator()

    try:
        contractor = Contractor(100, '/root/worker/demo/worker_demo.sh', '/dev/null')
        contractor.recruit()
        contractor.trial()

        while True:
            for i in xrange(contractor.idle_count()):
                id = IDGenerator.next()
                task = 'Task : %d' % id

                contractor.assign_task(id, task)

            contractor.handle_resign()
            contractor.recruit()

            results = contractor.wait_response(-1)
            print results
    except KeyboardInterrupt:
        contractor.quit()

